perf(moe): sync SM12x NVFP4 fused-MoE kernels to b12x HEAD - #4285
Conversation
Launch the resident-grid-barrier kernels cooperatively, keep the CTA-leader predicate in DSL IR, and restrict the fast FP4 quantizer to gated activations so relu2 keeps the exact path. Publish materialized tasks as expert/valid-rows pairs and decode tile/slice coordinates from the slot, dropping the ready-queue regime and five task-queue buffers. Size the dynamic kernel's pipeline stages from the real per-stage smem footprint, vectorize the static kernel's scatter epilogue, refresh the tile and MAC selection heuristics, bound the micro backend to decode-sized token counts, and reject workspaces past the 2^31 memref limit. Drop the unwired direct micro kernel and add clear_sm120_moe_caches(). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughSM12x MoE adds a direct-micro FP4 backend, compact dynamic task metadata, activation normalization, cooperative launches, tile-aware workspace caching, FP4 conversion helpers, and regression tests for gated activations and ReLU2. ChangesSM12x MoE execution
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant MoE_Dispatch
participant Dynamic_Workspace
participant MoEDynamicKernel
participant Task_Queue
participant MMA_DMA_Warps
MoE_Dispatch->>Dynamic_Workspace: select tile geometry and allocate compact buffers
MoE_Dispatch->>MoEDynamicKernel: launch compact task metadata
MoEDynamicKernel->>Task_Queue: publish and finalize deferred work items
MMA_DMA_Warps->>Task_Queue: claim and decode work items
Task_Queue->>MMA_DMA_Warps: provide expert, tile, slice, and valid-row data
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
| | `FLASHINFER_ROUTING_FORCE_BLOCK_PER_TOKEN` | unset | `csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu` | Forces the TRT-LLM MoE custom-routing kernel into "one-block-per-token" mode regardless of the active routing policy. Mainly used to reproduce specific perf points. | | ||
| | `FLASHINFER_B12X_MICRO_SHARE_INPUT` | `1` | `flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py` | `0` disables the B12x MoE micro-batch input-sharing optimization. Internal/experimental — leave at the default unless investigating an SM12x MoE regression. | | ||
| | `FLASHINFER_B12X_FORCE_MOE_W4A16` | unset | `flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py` | When set (any non-empty value), forces the SM12x MoE dispatcher onto the W4A16 kernel path regardless of weight dtype. Internal/experimental — used to reproduce W4A16-specific issues. | | ||
| | `FLASHINFER_B12X_MOE_TILE_MN` | unset | `flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py` | Overrides the SM12x MoE decode tile (`64x128` or `128x128`), bypassing the routed-rows heuristic. Applies to the micro kernel and the multi-top-k static kernel only. Internal/experimental — benchmarking aid. | |
There was a problem hiding this comment.
Can we refrain from adding a new environment variable FLASHINFER_B12X_MOE_TILE_MN?
I don't foresee any user using it so we can likely delete
| override = os.environ.get("FLASHINFER_B12X_MOE_TILE_MN") | ||
| if override: | ||
| if override not in ("64x128", "128x128"): | ||
| raise ValueError( | ||
| f"FLASHINFER_B12X_MOE_TILE_MN={override!r} is not supported. " | ||
| "The micro and static kernels only take 64x128 or 128x128." | ||
| ) | ||
| tile_m, tile_n = override.split("x") | ||
| return (int(tile_m), int(tile_n)) |
There was a problem hiding this comment.
Let's delete the environment variable here. I don't think any customer will be using this one
There was a problem hiding this comment.
Agreed, it was carried over from upstream as a benchmarking aid but has no real user. Removed in bf70250, along with the CLAUDE.md entry.
AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
… b12x MoE routine AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/bot run tests/moe |
Port upstream's current MoEDirectMicroKernel, which reads the top-k ids directly instead of using a routing pre-pass, add gelu_tanh to it, and route the smallest decode batches to it ahead of the MMA micro kernel. Let the dynamic kernel run M-tiles of 16, 32, or 64 rows picked from routed rows per expert, with the workspace cache keyed on the selected tile so cached workspaces keep the right geometry. Backend cutovers are provisional pending cross-GPU measurement. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py (1)
4899-4924: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClamp the cooperative micro-kernel grid by driver occupancy.
configurecapsgrid_xwithmin(get_num_sm(device), get_max_active_clusters(1)), butconfiguredoes not run occupancy analysis. The fused launch usescooperative=Truewithmin_blocks_per_mp=1, so any compiled micro-kernel with fewer than one CTA per SM can deadlock or fail. Clampgrid_xwith the occupancy returned for the compiled kernel, or release cooperative admission when occupancy is too low.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 4899 - 4924, Update the fused launch around the kernel’s .launch call to account for compiled-kernel driver occupancy before using cooperative=True. Clamp grid_x to the occupancy-supported CTA count for the compiled micro-kernel, or disable cooperative admission when occupancy is below one CTA per SM, while preserving the existing phase-specific barrier behavior and launch configuration.flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py (1)
1417-1421: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winResize
task_expert/task_valid_rowsby the materialized task count.
task_tailstoresexpert_tile_base[num_experts] * materialized_num_groups, wherematerialized_num_groupsdepends on_TASK_SLICE_CHUNK, whileallocate_sm120_dynamic_workspace()still usesmax_m_tiles * (_DYNAMIC_SLICE_CHUNK-rounded gate tile groups). With_TASK_SLICE_CHUNK = 1, every physical tile can materializegate_tile_cnttasks, sotask_expertandtask_valid_rowscan be undersized and allow out-of-range task metadata writes/consumption.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py` around lines 1417 - 1421, Resize the task_expert and task_valid_rows workspace allocations in allocate_sm120_dynamic_workspace() using the maximum materialized task count, matching the task_tail calculation of expert_tile_base[num_experts] * materialized_num_groups in the dynamic kernel. Account for _TASK_SLICE_CHUNK so the allocation remains sufficient when each physical tile materializes multiple gate-tile tasks, while preserving the existing metadata layout.
🧹 Nitpick comments (5)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py (2)
19-33: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake
is_gated_activationreuseGATED_MOE_ACTIVATIONS.
GATED_MOE_ACTIVATIONSat Line 22 and the tuple inis_gated_activationat Line 86 list the same activations. Two lists can drift when a new gated activation is added.moe_dispatch.pyimportsis_gated_activation, and the direct-micro kernel importsis_gated_moe_activation, so both must agree.♻️ Proposed consolidation
def is_gated_activation(activation: str) -> bool: - return activation in ("silu", "gelu_tanh", "swigluoai_uninterleave") + return activation in GATED_MOE_ACTIVATIONS🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py` around lines 19 - 33, Update the existing is_gated_activation helper to determine membership using the shared GATED_MOE_ACTIVATIONS constant instead of its duplicated activation tuple. Preserve the helper’s current input normalization and boolean behavior, and keep is_gated_moe_activation aligned with the same shared set.
68-82: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAlign the beta default with the rejection rule.
SWIGLUOAI_DEFAULT_BETAis 1.0, and the public MoE APIs also defaultswiglu_beta=1.0. For any activation other thanswigluoai_uninterleave, this function rejects a beta of 1.0 and only accepts 0.0._get_direct_micro_kernelavoids the error because it setsswiglu_beta = Nonefirst, but a direct call tobuild_direct_micro_kernel(..., activation="silu", swiglu_beta=1.0)raises. Accept the documented default value for non-gated-configurable activations, or document the required 0.0.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py` around lines 68 - 82, Update normalize_swiglu_beta_for_activation so the documented default swiglu_beta=1.0 is accepted for activations other than SWIGLUOAI_UNINTERLEAVE, while preserving rejection of other nonzero explicitly configured values and returning the inactive beta value used by those activations. Keep the configurable activation’s finite-beta validation unchanged.flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py (1)
5129-5148: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSilence the blind-except lint or narrow the caught type.
Ruff reports BLE001 for
except Exception as exc. The broad catch is intentional here, because the probe must never disable the process on DSL internals changes without a warning. Add an explicit suppression so the lint stays clean and the intent stays documented.♻️ Proposed suppression
- except Exception as exc: + except Exception as exc: # noqa: BLE001 - probe must degrade, never raise🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 5129 - 5148, Add an explicit Ruff BLE001 suppression to the intentional broad `except Exception as exc` handler in `compiled_direct_micro_accepts_block_dim`, preserving its warning-once behavior and safe `False` fallback. Keep the existing exception handling and explanatory comments unchanged.Source: Linters/SAST tools
tests/moe/test_b12x_fused_moe.py (1)
1865-1944: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
swigluoai_uninterleaveto the forced direct-micro activation coverage.This test parametrizes
siluandgelu_tanhonly. The direct-micro kernel builder special-casesSWIGLUOAI_UNINTERLEAVEexplicitly (it is the only activation for whichswiglu_alpha/swiglu_beta/swiglu_limitare kept instead of reset toNone). Since this test's stated purpose is to force the direct-micro backend and fail loudly instead of silently falling back, add at least oneswigluoai_uninterleavecase to confirm the direct-micro path handles its distinct swiglu parameter handling correctly, not just infer it from natural backend routing intest_activation_accuracy.✅ Suggested parametrize addition
`@pytest.mark.parametrize`( "activation,num_tokens,top_k", [ ("silu", 1, 2), ("silu", 2, 2), ("silu", 8, 2), ("silu", 1, 8), ("silu", 2, 8), ("silu", 8, 8), ("gelu_tanh", 1, 2), ("gelu_tanh", 8, 2), + ("swigluoai_uninterleave", 1, 2), + ("swigluoai_uninterleave", 8, 2), ], ) def test_direct_micro_forced_accuracy( self, monkeypatch, activation: str, num_tokens: int, top_k: int ):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_b12x_fused_moe.py` around lines 1865 - 1944, Add at least one “swigluoai_uninterleave” case to the activation parameter list for test_direct_micro_forced_accuracy, using a supported tiny decode shape, so the forced direct-micro path exercises its distinct SwiGLU parameter handling alongside silu and gelu_tanh.flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py (1)
2232-2259: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHoist the FC2 sub-tile re-slicing out of the per-slice loop.
This block re-slices
sA/sSFAat a fixed offset (Int32(0)), and that offset does not depend onslice_idx. The enclosingwhile slice_idx < task_slice_count_val:loop is a runtime device-side loop (its trip count comes fromwork_item[_WORK_SLICE_COUNT], a value known only at kernel run time), so this recomputation ofcute.local_tile,partition_A,_partition_fragment_SFA, and the retiles executes on every slice iteration, not once per task. Compare this to the FC1 case (lines 1742-1766), where the equivalent per-task re-slice is computed once, before the slice loop starts. For tasks withtask_slice_count_val > 1, hoisting the FC2 re-slice the same way removes redundant per-iteration work in this hot GEMM consumer loop.Introduce dedicated variable names for the FC2 views (for example
csA_fc2/crA_fc2/tCrA_fc2) computed once before the slice loop, sincecsA/crA/tCrAmust keep referring to the FC1 per-task view for the next slice iteration's Phase A.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py` around lines 2232 - 2259, Hoist the FC2 sub-tile setup currently inside the phase-B block out of the runtime slice loop, computing the sA/sSFA local tiles, partitions, fragments, and retiles once per task before the loop. Use dedicated FC2 variables such as csA_fc2, crA_fc2, tCrA_fc2 and corresponding SFA names, while preserving csA/crA/tCrA for the FC1 per-slice Phase A view; update Phase B to use the hoisted FC2 views.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 1-10: Update the module docstring for MoEDirectMicroKernel to
describe FC2 as writing bf16x2 results with plain indexed stores into the
token-major output, removing the atomic-scatter wording. Preserve the existing
description of token-major output and avoid implying accumulation semantics.
---
Outside diff comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 4899-4924: Update the fused launch around the kernel’s .launch
call to account for compiled-kernel driver occupancy before using
cooperative=True. Clamp grid_x to the occupancy-supported CTA count for the
compiled micro-kernel, or disable cooperative admission when occupancy is below
one CTA per SM, while preserving the existing phase-specific barrier behavior
and launch configuration.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py`:
- Around line 1417-1421: Resize the task_expert and task_valid_rows workspace
allocations in allocate_sm120_dynamic_workspace() using the maximum materialized
task count, matching the task_tail calculation of expert_tile_base[num_experts]
* materialized_num_groups in the dynamic kernel. Account for _TASK_SLICE_CHUNK
so the allocation remains sufficient when each physical tile materializes
multiple gate-tile tasks, while preserving the existing metadata layout.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py`:
- Around line 19-33: Update the existing is_gated_activation helper to determine
membership using the shared GATED_MOE_ACTIVATIONS constant instead of its
duplicated activation tuple. Preserve the helper’s current input normalization
and boolean behavior, and keep is_gated_moe_activation aligned with the same
shared set.
- Around line 68-82: Update normalize_swiglu_beta_for_activation so the
documented default swiglu_beta=1.0 is accepted for activations other than
SWIGLUOAI_UNINTERLEAVE, while preserving rejection of other nonzero explicitly
configured values and returning the inactive beta value used by those
activations. Keep the configurable activation’s finite-beta validation
unchanged.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 5129-5148: Add an explicit Ruff BLE001 suppression to the
intentional broad `except Exception as exc` handler in
`compiled_direct_micro_accepts_block_dim`, preserving its warning-once behavior
and safe `False` fallback. Keep the existing exception handling and explanatory
comments unchanged.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py`:
- Around line 2232-2259: Hoist the FC2 sub-tile setup currently inside the
phase-B block out of the runtime slice loop, computing the sA/sSFA local tiles,
partitions, fragments, and retiles once per task before the loop. Use dedicated
FC2 variables such as csA_fc2, crA_fc2, tCrA_fc2 and corresponding SFA names,
while preserving csA/crA/tCrA for the FC1 per-slice Phase A view; update Phase B
to use the hoisted FC2 views.
In `@tests/moe/test_b12x_fused_moe.py`:
- Around line 1865-1944: Add at least one “swigluoai_uninterleave” case to the
activation parameter list for test_direct_micro_forced_accuracy, using a
supported tiny decode shape, so the forced direct-micro path exercises its
distinct SwiGLU parameter handling alongside silu and gelu_tanh.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 014f8ce1-2422-46cc-bb24-d6ed8b33299f
📒 Files selected for processing (7)
flashinfer/cute_dsl/fp4_common.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/__init__.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.pytests/moe/test_b12x_fused_moe.py
AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py (3)
656-699: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winValidate A8-MX paired-block K before launch.
A8-MX reads
pair_delta = ±_BLOCK_SIZEoffsets for every_BLOCK_SIZE-sized FC1/FC2 block, including the final block, while the A8-MX scale column is derived fromcfg.k_dim // 32. Ifcfg.k_dimis an odd multiple of_BLOCK_SIZE, the final FC1 paired read can go pasta_inputand the scale layout does not contain the full pair. Reject incomplete K pairs at configure time.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 656 - 699, In the A8-MX configuration path, validate that cfg.k_dim contains complete paired _BLOCK_SIZE blocks before launch, rejecting odd multiples that leave the final FC1/FC2 block without its partner. Add this check alongside the existing a8_mx_mode validation near the cfg.i_chunk % 32 guard, and raise ValueError for incomplete K pairs.
3962-4049: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAdd E8M0 loading to the FC1 tail and aligned scale paths.
_ld_e8m0_scaleis only called from explicitly gated K-segment branches; the FC1 tail and genericcfg.k_segments_alignedpaths still select only W4A16 E4M3 or legacy E4M3 loaders. Addscale_format_e8m0_k32E8M0 loading here, or reject E8M0 for these paths, to prevent FC1 scale decoding from using the wrong layout.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 3962 - 4049, Update the FC1 scale-loading logic around the w4a16 tail branch and the generic cfg.k_segments_aligned path to handle scale_format_e8m0_k32 explicitly. Reuse _ld_e8m0_scale with the correct K32 addressing and lane/segment indexing before selecting E4M3 loaders, or explicitly reject E8M0 for these paths; do not decode E8M0 through the legacy E4M3 layout.
4511-4601: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject or implement dynamic down-scale for W4A16 and A8-MX.
dynamic_down_scaleenablesfc2_rescale, but only the non-W4A16/non-A8-MX FC2 branch multiplies the packed outputs by it. The W4A16 branch packsv0, v1directly, and the A8-MX branch only foldsgs_fc2into the dequant values. Add constructor validation or apply the rescale consistently in those outputs.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 4511 - 4601, Ensure dynamic_down_scale is either rejected for W4A16 and A8-MX configurations through constructor validation, or applied consistently to their FC2 outputs. Update the A8-MX packing in the dynamic_down_scale path and the W4A16 packing around fc2_rescale so both use the computed rescale, matching the existing non-W4A16/non-A8-MX branch; otherwise add validation in the relevant constructor to disallow these combinations.
🧹 Nitpick comments (1)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py (1)
332-378: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueMake
fast_matha supported no-op or remove it from the direct micro API.
moedirectdispatch passesfast_maththroughbuild_direct_micro_kernel, butMoEDirectMicroKerneldiscards it immediately. Add back the ignored-argument handling and cache/no-op handling so the API contract is explicit; otherwise removefast_math, which changes the exported direct micro API.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 332 - 378, The MoEDirectMicroKernel API currently accepts fast_math but silently discards it without explicit no-op handling. Update MoEDirectMicroKernel to preserve fast_math as a documented supported no-op, including the expected ignored-argument/cache handling used by build_direct_micro_kernel, or remove fast_math consistently from the direct micro API and its dispatch path.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 656-699: In the A8-MX configuration path, validate that cfg.k_dim
contains complete paired _BLOCK_SIZE blocks before launch, rejecting odd
multiples that leave the final FC1/FC2 block without its partner. Add this check
alongside the existing a8_mx_mode validation near the cfg.i_chunk % 32 guard,
and raise ValueError for incomplete K pairs.
- Around line 3962-4049: Update the FC1 scale-loading logic around the w4a16
tail branch and the generic cfg.k_segments_aligned path to handle
scale_format_e8m0_k32 explicitly. Reuse _ld_e8m0_scale with the correct K32
addressing and lane/segment indexing before selecting E4M3 loaders, or
explicitly reject E8M0 for these paths; do not decode E8M0 through the legacy
E4M3 layout.
- Around line 4511-4601: Ensure dynamic_down_scale is either rejected for W4A16
and A8-MX configurations through constructor validation, or applied consistently
to their FC2 outputs. Update the A8-MX packing in the dynamic_down_scale path
and the W4A16 packing around fc2_rescale so both use the computed rescale,
matching the existing non-W4A16/non-A8-MX branch; otherwise add validation in
the relevant constructor to disallow these combinations.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 332-378: The MoEDirectMicroKernel API currently accepts fast_math
but silently discards it without explicit no-op handling. Update
MoEDirectMicroKernel to preserve fast_math as a documented supported no-op,
including the expected ignored-argument/cache handling used by
build_direct_micro_kernel, or remove fast_math consistently from the direct
micro API and its dispatch path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5d13c5f2-9adc-4011-9794-76dd28cc166d
📒 Files selected for processing (4)
flashinfer/cute_dsl/fp4_common.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py
🚧 Files skipped from review as they are similar to previous changes (3)
- flashinfer/cute_dsl/fp4_common.py
- flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py
- flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
Gate the direct micro kernel on small intermediate sizes and fewer routed rows, where its CUDA-core dots win. Tie the dispatch task-queue sizing to the kernel's materialization granularity, accept the module defaults for swiglu alpha and beta on activations that ignore them, derive is_gated_activation from the gated set, document fast_math as a no-op, and cover swigluoai in the forced direct-micro tests. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py (4)
792-795: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMask invalid intermediate lanes in both narrow FC2 paths.
When
w_valid == 0, these paths still readintermediate. The wide paths already mask the same reads because zero weights do not neutralize NaN values from uninitialized storage. Apply the same mask toxh0throughxh3here.Proposed fix
- xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + xh0 = ( + Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + )Apply the same pattern to
xh1,xh2, andxh3in both functions.Also applies to: 1222-1226
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 792 - 795, Update the narrow FC2 paths containing the xh0–xh3 loads in both functions to mask each intermediate read with w_valid, matching the existing wide-path behavior. Apply the mask consistently to xh0, xh1, xh2, and xh3 so invalid lanes do not consume intermediate values.
460-470: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReject packed scale layouts with a non-64-aligned FC1 row width.
Both permutations operate within 64-element groups.
is_supporteddoes not enforce this precondition. The FC1 scale layouts allocate exactlycfg.two_ncolumns. For example,_packed_scale_col(4)returns32; that is outside a row with widthcfg.two_n == 32.Add a configure-time check for
cfg.two_n % 64 == 0when W4A16 or packed E8M0 scales are used. Otherwise, use a padded or logical layout.Proposed validation
+ packed_fc1_scales = self.w4a16_mode or ( + self.scale_format_e8m0_k32 + and not self.e8m0_scale_layout_logical + ) + if packed_fc1_scales and cfg.two_n % 64 != 0: + raise ValueError("packed FC1 scales require cfg.two_n % 64 == 0")Also applies to: 493-503, 597-627, 4842-4850
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 460 - 470, Add configure-time validation in the support/configuration path covering W4A16 and packed E8M0 scale layouts, requiring cfg.two_n to be divisible by 64 before constructing the FC1 scale layouts or using _packed_scale_col. Reject unsupported configurations clearly; do not allow the existing permutations to index beyond the allocated cfg.two_n columns.
4546-4562: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPrevent A8-MX FC2 from reading past the intermediate tile.
pair_deltapairs each 32-element block with an adjacent block. The current validation only requirescfg.i_chunk % 32 == 0. If the tile contains an odd number of 32-element blocks, the final even block uses+_BLOCK_SIZEand reads beyondsmem_int.Require an even number of blocks, or make the final block pair with itself.
Proposed validation
- if self.a8_mx_mode and self.compile_time_phase != 1 and cfg.i_chunk % 32 != 0: + if ( + self.a8_mx_mode + and self.compile_time_phase != 1 + and cfg.i_chunk % (2 * _BLOCK_SIZE) != 0 + ):Also applies to: 698-701
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 4546 - 4562, Update the A8-MX FC2 validation used by the micro-kernel setup and the corresponding path near the block-processing logic to require cfg.i_chunk to contain an even number of _BLOCK_SIZE elements, or otherwise make the final unpaired block use itself as its partner. Ensure pair_delta in the a8_mx_mode branch cannot index beyond smem_int while preserving adjacent pairing for all complete block pairs.
1238-1373: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse E8M0 addressing in
_m2_fc2_rowquad_narrow.When
scale_format_e8m0_k32is enabled, this function still uses the legacyebase_sfandbsf_offcalculations. The E8M0 scale tensor uses the[E, N/32, K]layout. Converting the loaded byte withcvt_e8m0_to_f32does not correct the address calculation.Add the
_ld_e8m0_scalebranch used by the other FC2 paths forbsf_f0throughbsf_f3.Required change
+ if cutlass.const_expr(self.scale_format_e8m0_k32): + ebase_w2p = Int64(eid) * Int64((cfg.n // 32) * cfg.k_dim) + kb32_i = lane >> Int32(2) + bsf_f0 = self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row0, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + elif cutlass.const_expr(self.w4a16_mode): + ... + else: + ...Repeat the format-specific addressing for rows 1 through 3.
Also applies to: 4868-4872
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py` around lines 1238 - 1373, Update _m2_fc2_rowquad_narrow so bsf_f0 through bsf_f3 use the existing _ld_e8m0_scale branch when scale_format_e8m0_k32 is enabled. Use E8M0 addressing for the [E, N/32, K] layout, including the corresponding row-specific arguments for rows 1–3, while preserving the current legacy addressing path for other scale formats.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py`:
- Around line 792-795: Update the narrow FC2 paths containing the xh0–xh3 loads
in both functions to mask each intermediate read with w_valid, matching the
existing wide-path behavior. Apply the mask consistently to xh0, xh1, xh2, and
xh3 so invalid lanes do not consume intermediate values.
- Around line 460-470: Add configure-time validation in the
support/configuration path covering W4A16 and packed E8M0 scale layouts,
requiring cfg.two_n to be divisible by 64 before constructing the FC1 scale
layouts or using _packed_scale_col. Reject unsupported configurations clearly;
do not allow the existing permutations to index beyond the allocated cfg.two_n
columns.
- Around line 4546-4562: Update the A8-MX FC2 validation used by the
micro-kernel setup and the corresponding path near the block-processing logic to
require cfg.i_chunk to contain an even number of _BLOCK_SIZE elements, or
otherwise make the final unpaired block use itself as its partner. Ensure
pair_delta in the a8_mx_mode branch cannot index beyond smem_int while
preserving adjacent pairing for all complete block pairs.
- Around line 1238-1373: Update _m2_fc2_rowquad_narrow so bsf_f0 through bsf_f3
use the existing _ld_e8m0_scale branch when scale_format_e8m0_k32 is enabled.
Use E8M0 addressing for the [E, N/32, K] layout, including the corresponding
row-specific arguments for rows 1–3, while preserving the current legacy
addressing path for other scale formats.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 0cb8f5db-3185-45ed-a565-9e0694ecfce6
📒 Files selected for processing (4)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.pytests/moe/test_b12x_fused_moe.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/moe/test_b12x_fused_moe.py
- flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py
- flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
|
/bot run tests/moe |
|
[FAILED] Pipeline #60494924 — 4/18 executed test jobs passed Compared with nightly #60276939. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
Could not compare
|
|
/bot run tests/moe |
|
[FAILED] Pipeline #60881505 — 16/18 executed test jobs passed Compared with nightly #60712014. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
No individual test or infrastructure failures could be extracted. |
## 📌 Description Syncs the SM12x NVFP4 (W4A4) fused-MoE kernels behind `b12x_fused_moe` and `B12xMoEWrapper` with upstream [b12x](https://github.com/local-inference-lab/sparkinfer) at `f9be272`, reaching backend parity. The module now has four kernels. `MoEDirectMicroKernel` takes the smallest decode batches, `MoEMicroKernel` and `MoEStaticKernel` cover the rest of decode with tile assignments fixed at launch, and `MoEDynamicKernel` covers prefill by distributing tiles through an in-kernel work queue. Changes: - Adopt upstream's current `MoEDirectMicroKernel`, replacing a stale copy in `moe_direct_micro_kernel.py` that was never routed to. It reads the top-k ids directly, so the smallest batches skip the routing pre-pass that dominates their runtime, and it now supports gelu_tanh. - Let `MoEDynamicKernel` use tiles smaller than 128 rows, so sparse routing stops padding every expert up to 128. - Launch all kernels cooperatively. They synchronize the whole grid between phases, which deadlocks if concurrent work keeps part of the grid off the GPU. - Simplify `MoEDynamicKernel`'s work queue and size its pipeline stages from the real shared-memory footprint. - Vectorize `MoEStaticKernel`'s scatter epilogue and refresh the decode tile heuristics. - Restrict the fast FP4 quantizer to gated activations, since relu2's squared outputs need the exact one. Public API and behavior changes: none. ## 📊 Performance Versus main, measured over 38 target-model MoE shapes on three GPUs. Geomean speedup per shape group: | shape group | DGX Spark | RTX 5080 | RTX Pro 6000 SE | |---|---|---|---| | decode, batch 1-2 | 1.06x | 1.30x | 1.40x | | decode, batch 4-40 | 1.00x | 1.00x | 1.00x | | prefill, batch 512-2048 | 1.10x | 1.10x | 1.09x | | prefill, batch 8192 | 1.00x | 1.00x | 1.00x | - The batch 1-2 gains come from the new `MoEDirectMicroKernel`; the shapes it serves win up to 2.3x. - The prefill gains come from `MoEDynamicKernel` picking tiles smaller than 128 rows; those shapes win up to 1.26x. ## 🔍 Related Issues #4223 (item 2). ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - On SM121 (DGX Spark GB10): `tests/moe/test_b12x_fused_moe.py` (157 cases) and `tests/moe/test_unified_moe_b12x.py` (58 cases) pass, including tests that force each backend and a test that a smaller tile is picked again after a large-batch call. ## Reviewer Notes Deviations from the b12x source, beyond import plumbing: - gelu_tanh added to `MoEDirectMicroKernel`, matching the activation coverage of the module's other kernels. - The module's API accepts reciprocal-form scales but upstream's kernel only takes multipliers, so dispatch inverts them before launch. - `MoEDirectMicroKernel`'s cutover is re-measured: upstream's only alternative is its counterpart of `MoEDynamicKernel`, while here it competes with `MoEMicroKernel` and `MoEStaticKernel`. - Upstream sizes scratch and builds the kernel together inside one plan object, so their tile sizes always match. FlashInfer has no plan object, so the workspace cache is keyed by tile size to give the same guarantee. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com> (cherry picked from commit 5a4b54b)
4 participants
📌 Description
Syncs the SM12x NVFP4 (W4A4) fused-MoE kernels behind
b12x_fused_moeandB12xMoEWrapperwith upstream b12x atf9be272, reaching backend parity.The module now has four kernels.
MoEDirectMicroKerneltakes the smallest decode batches,MoEMicroKernelandMoEStaticKernelcover the rest of decode with tile assignments fixed at launch, andMoEDynamicKernelcovers prefill by distributing tiles through an in-kernel work queue.Changes:
MoEDirectMicroKernel, replacing a stale copy inmoe_direct_micro_kernel.pythat was never routed to. It reads the top-k ids directly, so the smallest batches skip the routing pre-pass that dominates their runtime, and it now supports gelu_tanh.MoEDynamicKerneluse tiles smaller than 128 rows, so sparse routing stops padding every expert up to 128.MoEDynamicKernel's work queue and size its pipeline stages from the real shared-memory footprint.MoEStaticKernel's scatter epilogue and refresh the decode tile heuristics.Public API and behavior changes: none.
📊 Performance
Versus main, measured over 38 target-model MoE shapes on three GPUs. Geomean speedup per shape group:
MoEDirectMicroKernel; the shapes it serves win up to 2.3x.MoEDynamicKernelpicking tiles smaller than 128 rows; those shapes win up to 1.26x.🔍 Related Issues
#4223 (item 2).
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
tests/moe/test_b12x_fused_moe.py(157 cases) andtests/moe/test_unified_moe_b12x.py(58 cases) pass, including tests that force each backend and a test that a smaller tile is picked again after a large-batch call.Reviewer Notes
Deviations from the b12x source, beyond import plumbing:
MoEDirectMicroKernel, matching the activation coverage of the module's other kernels.MoEDirectMicroKernel's cutover is re-measured: upstream's only alternative is its counterpart ofMoEDynamicKernel, while here it competes withMoEMicroKernelandMoEStaticKernel.🤖 Generated with Claude Code